# 设计思路：
# 1、定义超参数，提高模型的稳定性和效率
# 2、定义数据转换，对原始数据进行预处理
# 3、加载数据集
# 4、定义卷积神经网络模型
# 5、创建模型，定义损失函数和优化器
# 6、训练模型
# 7、测试模型性能

# 导入库
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 设置超参数
batch_size = 16
learning_rate = 0.001
num_epochs = 10

# 数据转换
# Resize将图像大小调整为 224×224；ToTensor将图像转换为张量；Normalize对图像进行标准化处理，使其均值为 0.5，方差为 0.5。
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

# 加载数据集
# ImageFolder直接读取文件夹中的图片数据集，并按照文件夹名字进行分类。
train_dataset = datasets.ImageFolder('D:\software\PyCharm Community Edition 2022.3.2\jbr\\bin\D：、python-learning\weather_dataset\\train', transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

test_dataset = datasets.ImageFolder('D:\software\PyCharm Community Edition 2022.3.2\jbr\\bin\D：、python-learning\weather_dataset\\test', transform=transform)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

# 定义卷积神经网络模型
class CNN(nn.Module):
    # 初始化模型
    def __init__(self):
        super(CNN, self).__init__()
        self.conv_layer1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.pooling1 = nn.MaxPool2d(kernel_size=2)
        self.conv_layer2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.pooling2 = nn.MaxPool2d(kernel_size=2)
        self.conv_layer3 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(64*56*56, 1024)
        self.fc2 = nn.Linear(1024, 11)
        # 其中，使用了三个卷积层，两个最大池化层和两个全连接层。最后一层的输出大小为11，因为一共有11种天气。

    # 前馈函数
    def forward(self, x):
        x = self.pooling1(torch.relu(self.conv_layer1(x)))
        x = self.pooling2(torch.relu(self.conv_layer2(x)))
        # 将多维张量展平为一维张量以符合全连接层输入
        x = self.flatten(torch.relu(self.conv_layer3(x)))
        x = self.fc2(torch.relu(self.fc1(x)))
        return x
                                                                                          
# 创建模型
cnn = CNN()
# 分类问题，定义损失函数为交叉熵
criterion = nn.CrossEntropyLoss()
# 经多次测试，定义优化器为Adagrad优化器，准确率较高
optimizer = optim.Adagrad(cnn.parameters(), lr=learning_rate)

# 训练模型
for epoch in range(num_epochs):
    for images, labels in train_loader:
        # 每次循环前使用optimizer.zero_grad()方法来将梯度清零，以防止梯度累积影响模型性能
        optimizer.zero_grad()
        # 将训练集图像传入cnn模型，经处理后得到输出值为outputs
        outputs = cnn(images)
        # 将模型的输出和标签传入交叉熵损失函数，得到损失值（“错误程度”）
        loss = criterion(outputs, labels)
        # 从损失值节点开始，通过反向传播，计算loss关于所有变量的导数，并将结果存在对应变量的 .grad 属性中
        loss.backward()
        # 根据这些梯度(grad)对模型参数进行更新，使得模型在下一轮迭代中更加接近最优解
        optimizer.step()
    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))
# 每个epoch中遍历一次训练集，并计算损失和梯度，进行反向传播更新参数。

# 测试模型性能
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        outputs = cnn(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        accuracy = 100 * correct / total

print(f'模型在测试集上的准确率为: {accuracy} %')